import numpy as np
import faiss
import time
import logging

# ---------参数设置------------------
logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s:%(message)s", level=logging.DEBUG)
logger = logging.getLogger(__name__)
k = 20
# --------特征池的获取-------------
start_ = time.time()
all_feature= np.load('all_feature_11000.npy').astype('float32')
all_id = np.load('all_id_11000.npy')
end_ = time.time()
logger.debug('特征池建立时间： ' + str(end_ - start_))

# ------------目标特征的提取--------
obj_feature = np.load('video_obj_pingpang.npy').astype('float32')
end_ = time.time()-end_
logger.debug('目标建立时间： ' + str(end_ ))
# -------------相似向量检测---------
'''精确方法-brute-force'''
# IndexFlatIP  为 cos  L2为蛮力匹配
def cos_dis():
    logger.debug('开始精确方法-brute-force 进行特征向量匹配：')
    start_ = time.time()
    index = faiss.IndexFlatL2(2048)
    # assert not index.is_trained
    # print(index.ntotal)
    index.add(all_feature)
    # print(index.ntotal)
    # 返回结果数

    dis,ind = index.search(obj_feature,k)
    end_ = time.time()
    out_put(dis, ind)
    logger.debug('排序时间(brute)： '+str(end_-start_))
    print('\n')

'''倒排表快速索引'''
def inverted_index():
    logger.debug('开始较快方法-倒排表 进行特征向量匹配：')
    start_ = time.time()
    nlist = 300 #特征池分割为50块

    quantizer = faiss.IndexFlatL2(2048)  # 量化器
    index = faiss.IndexIVFFlat(quantizer, 2048, nlist, faiss.METRIC_L2)
           # METRIC_L2计算L2距离, 或faiss.METRIC_INNER_PRODUCT计算内积
    end_quantizer = time.time()
    logger.debug('quantizer时间：'+str(end_quantizer-start_))
    assert not index.is_trained   #倒排表索引类型需要训练
    index.train(all_feature)  # 训练数据集应该与数据库数据集同分布
    assert index.is_trained
    end_train = time.time()
    logger.debug('train时间：' + str(end_train - end_quantizer))
    index.add(all_feature)
    end_add = time.time()
    logger.debug('add时间：' + str(end_add - end_train))
    index.nprobe = 5  # 选择n个维进行索引,
    dis, ind = index.search(obj_feature,k)
    end_index = time.time()
    logger.debug('index时间：' + str(end_index - end_add))
    end_ = time.time()
    out_put(dis,ind)
    logger.debug('排序时间(inverted)： '+str(end_-start_))
    print('\n')

'''乘积量化索引'''
def PQ_index():
    logger.debug('开始量化方法-PQ 进行特征向量匹配：')
    start_ = time.time()
    # 乘积量化方法保存原始向量的有损压缩
    nlist = 300
    m = 32                             # 列方向划分个数，必须能被d整除

    quantizer = faiss.IndexFlatL2(2048)
    end_quantizer = time.time()
    logger.debug('quantizer时间：'+str(end_quantizer-start_))
    # 8 表示每个子向量被编码为 8 bits
    index = faiss.IndexIVFPQ(quantizer, 2048, nlist, m, 8)
    index.train(all_feature)
    end_train = time.time()
    logger.debug('train时间：' + str(end_train - end_quantizer))

    index.add(all_feature)
    end_add = time.time()
    logger.debug('add时间：' + str(end_add - end_train))
    index.nprobe = 5
    dis, ind = index.search(obj_feature,k)
    end_index = time.time()
    logger.debug('index时间：' + str(end_index - end_add))
    end_ = time.time()
    out_put(dis, ind)
    logger.debug('排序时间总时间(PQ)： ' + str(end_ - start_))
    print('\n')
    # faiss.index_factory(d, "PCAR16,IVF50,SQ8") 1.降维 2.聚类 3.scalar 量化，每个向量编码为8bit 不支持GPU
    # index = index_factory(128,"OPQ16_64,IMI2x8,PQ8+16")  # 原始向量128维，用OPQ降为64维，分为16类，用2*8bit的倒排多索引
    # # ，用PQ编码为8byte保存，检索时使用16byte。

'''GPU'''
def GGGG():
    logger.debug('开始倒排表基于GPU 进行特征向量匹配：')
    start_ = time.time()
    nlist = 300
    quantizer = faiss.IndexFlatL2(2048)
    end_quantizer = time.time()
    logger.debug('quantizer时间：'+str(end_quantizer-start_))
    index = faiss.IndexIVFFlat(quantizer,2048,nlist,faiss.METRIC_L2)
    gpu_index = faiss.index_cpu_to_all_gpus(index)
    end_gpu = time.time()
    logger.debug('cpu=>gpu时间：' + str(end_gpu - end_quantizer))
    gpu_index.train(all_feature)
    end_train = time.time()
    logger.debug('train时间：' + str(end_train - end_gpu))
    gpu_index.add(all_feature)
    end_add = time.time()
    logger.debug('add时间：'+str(end_add - end_train))
    gpu_index.nprobe = 5
    # 返回结果数

    dis,ind = gpu_index.search(obj_feature,k)
    end_index = time.time()
    logger.debug('index时间：'+str(end_index - end_add))
    end_ = time.time()
    out_put(dis, ind)
    logger.debug('排序时间(inverted_GPU)： '+str(end_-start_))
    print('\n')

def out_put(dis,ind):
    # 升序返回每个查询向量的距离
    print('目标向量和特征池向量的距离: ')
    print(dis)
    # 升序返回每个查询向量的k个相似结果
    print('目标向量的相似向量结果: ')
    print(ind)
    ans = []
    for i in range(0,len(ind)):
        for j in range(0,k):
            out_id =int(all_id[ind[i, j], 1])
            ans.append(out_id)

    counts = np.bincount(ans)
    for i in range(0, 8):
        ss = np.argmax(counts)
        if i == 0:
            print('最相似的视频ID为：' + str(ss))
        if i == 1:
            print('次相似的视频ID为：' + str(ss))
        if i > 1 and ss != 0:
            print('其余参考相似的视频ID为：' + str(ss))
        counts[ss] = 0



if __name__ == '__main__':

    cos_dis()
    inverted_index()
    PQ_index()
    GGGG()